import numpy as np
import coptpy as cp
from coptpy import COPT

from admm.ADMM_update import terminal_args

def ADMM_acline_test(
    acline_args, terminal1: terminal_args, terminal2: terminal_args, rou: float, T: int
)->dict:
    """
    交流线路，无决策变量
    只有状态变量：两端有功，电压幅值平方，电流幅值平方
    """
    # Create COPT environment
    env = cp.Envr()

    # Create COPT model
    model = env.createModel("device_acline")

    # Add variables
    Pi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pi",
    )
    Pj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pj",
    )
    Qi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qi",
    )
    Qj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qj",
    )
    l2 = model.addMVar(
        (T, 1), lb=0.0, ub=acline_args["I2max"], vtype=COPT.CONTINUOUS, nameprefix="I"
    )
    V2i = model.addMVar(
        (T, 1),
        lb=acline_args["V2min"],
        ub=acline_args["V2max"],
        vtype=COPT.CONTINUOUS,
        nameprefix="V2i",
    )
    V2j = model.addMVar(
        (T, 1),
        lb=acline_args["V2min"],
        ub=acline_args["V2max"],
        vtype=COPT.CONTINUOUS,
        nameprefix="V2j",
    )
    cnorm2 = model.addVar(lb=0.0, ub=COPT.INFINITY, name="cnorm2")
    half_l2 = model.addMVar(
        (T, 1),
        lb=0.0,
        ub=acline_args["I2max"] / 2.0,
        vtype=COPT.CONTINUOUS,
        nameprefix="halfI",
    )
    auxPi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="auxPi",
    )
    auxPj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="auxPj",
    )
    auxQi = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qi",
    )
    auxQj = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Qj",
    )
    auxV2i = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="V2i",
    )
    auxV2j = model.addMVar(
        (T, 1),
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="V2j",
    )
    # Add linear constraints
    model.addConstrs(Pi + Pj == l2 * acline_args["R"], nameprefix="CP")
    model.addConstrs(Qi + Qj == l2 * acline_args["X"], nameprefix="CQ")
    model.addConstrs(
        V2j
        == V2i
        - 2 * (Pi * acline_args["R"] + Qi * acline_args["X"])
        + l2 * (acline_args["R"] ** 2 + acline_args["X"] ** 2),
        nameprefix="CV",
    )
    # Add auxiliary varibles constraints
    model.addConstrs(half_l2 == l2 / 2.0, nameprefix="ChalfI")
    model.addConstrs(auxPi == Pi - terminal1.P, nameprefix="CauxPi")
    model.addConstrs(auxQi == Qi - terminal1.Q, nameprefix="CauxQi")
    model.addConstrs(auxV2i == V2i - terminal1.V2, nameprefix="CauxV2i")
    model.addConstrs(auxPj == Pj - terminal2.P, nameprefix="CauxPj")
    model.addConstrs(auxQj == Qj - terminal2.Q, nameprefix="CauxQj")
    model.addConstrs(auxV2j == V2j - terminal2.V2, nameprefix="CauxV2j")

    # Add 锥约束
    for k in range(0, T):
        model.addCone(
            [half_l2[k].item(), V2i[k].item(), Pi[k].item(), Qi[k].item()],
            COPT.CONE_RQUAD,
        )
    # 惩罚二次项就写成标准二阶锥
    CONE_cvars = [cnorm2]
    CONE_cvars.extend(auxPi.tolist())
    CONE_cvars.extend(auxQi.tolist())
    CONE_cvars.extend(auxV2i.tolist())
    CONE_cvars.extend(auxPj.tolist())
    CONE_cvars.extend(auxQj.tolist())
    CONE_cvars.extend(auxV2j.tolist())
    model.addCone(CONE_cvars, COPT.CONE_QUAD)

    # Set linear objective
    obj = (
        terminal1.LP.transpose() @ auxPi
        + terminal1.LQ.transpose() @ auxQi
        + terminal1.LV2.transpose() @ auxV2i
        + terminal2.LP.transpose() @ auxPj
        + terminal2.LQ.transpose() @ auxQj
        + terminal2.LV2.transpose() @ auxV2j
    )
    obj = obj + rou / 2.0 * cnorm2
    model.setParam(COPT.Param.Logging, 0)
    model.setObjective(obj, COPT.MINIMIZE)
    model.solve()

    return {
        "status": model.status,
        "objval": model.objval,
        "Pi": [(v.name, v.x) for v in Pi.tolist()],
        "Pj": [(v.name, v.x) for v in Pj.tolist()],
        "Pim": Pi.x.transpose(),
        "sol_type": type(Pi.x),
    }


def create_acline(R: float, X: float, Vmax: float, Vmin: float, Imax: float) -> dict:
    return {"R": R, "X": X, "V2max": Vmax**2, "V2min": Vmin**2, "I2max": Imax**2}


def start():
    T=36
    T1 = terminal_args(T)
    T1.P=np.random.random((T,1))
    T1.LP = np.random.random((T, 1))
    T1.Q = np.random.random((T, 1))
    T1.LQ = np.random.random((T, 1))
    T1.V2 = np.random.random((T, 1))
    T1.LV2 = np.random.random((T, 1))
    T2 = T1
    acl = create_acline(0.1, 1.0, 10.7, 9.0, 3)
    ret=ADMM_acline_test(acl,T1,T2,0.7,T)
    if ret['status'] != COPT.OPTIMAL:
        print('好像没解出来')
        return
